/**
 * @since 1.0.0
 */
import { prepareTools as prepareAnthropicTools } from "@effect/ai-anthropic/AnthropicLanguageModel"
import * as AnthropicTool from "@effect/ai-anthropic/AnthropicTool"
import * as AiError from "@effect/ai/AiError"
import type * as IdGenerator from "@effect/ai/IdGenerator"
import * as LanguageModel from "@effect/ai/LanguageModel"
import * as AiModel from "@effect/ai/Model"
import type * as Prompt from "@effect/ai/Prompt"
import type * as Response from "@effect/ai/Response"
import { addGenAIAnnotations } from "@effect/ai/Telemetry"
import * as Tool from "@effect/ai/Tool"
import * as Context from "effect/Context"
import * as DateTime from "effect/DateTime"
import * as Effect from "effect/Effect"
import * as Encoding from "effect/Encoding"
import { dual } from "effect/Function"
import * as Layer from "effect/Layer"
import * as Predicate from "effect/Predicate"
import * as Stream from "effect/Stream"
import type { Span } from "effect/Tracer"
import type { Mutable, Simplify } from "effect/Types"
import { AmazonBedrockClient } from "./AmazonBedrockClient.js"
import type {
  BedrockFoundationModelId,
  CachePointBlock,
  ContentBlock,
  ConverseRequest,
  ConverseResponse,
  ConverseResponseStreamEvent,
  ConverseTrace,
  DocumentFormat,
  Message,
  SystemContentBlock,
  Tool as AmazonBedrockTool,
  ToolChoice,
  ToolConfiguration
} from "./AmazonBedrockSchema.js"
import { ImageFormat } from "./AmazonBedrockSchema.js"
import * as InternalUtilities from "./internal/utilities.js"

const BEDROCK_CACHE_POINT: {
  readonly cachePoint: typeof CachePointBlock.Encoded
} = { cachePoint: { type: "default" } }

/**
 * @since 1.0.0
 * @category Models
 */
export type Model = typeof BedrockFoundationModelId.Encoded

// =============================================================================
// Configuration
// =============================================================================

/**
 * @since 1.0.0
 * @category Context
 */
export class Config extends Context.Tag(
  "@effect/ai-amazon-bedrock/AmazonBedrockLanguageModel/Config"
)<Config, Config.Service>() {
  /**
   * @since 1.0.0
   */
  static readonly getOrUndefined: Effect.Effect<typeof Config.Service | undefined> = Effect.map(
    Effect.context<never>(),
    (context) => context.unsafeMap.get(Config.key)
  )
}

/**
 * @since 1.0.0
 */
export declare namespace Config {
  /**
   * @since 1.0.0
   * @category Configuration
   */
  export interface Service extends
    Simplify<
      Partial<
        Omit<
          typeof ConverseRequest.Encoded,
          "messages" | "system" | "toolConfig"
        >
      >
    >
  {}
}

// =============================================================================
// Amazon Bedrock Provider Options / Metadata
// =============================================================================

/**
 * @since 1.0.0
 * @category Provider Options
 */
export type AmazonBedrockReasoningInfo = {
  readonly type: "thinking"
  /**
   * Thinking content as an encrypted string, which is used to verify
   * that thinking content was indeed generated by Amazon Bedrock's API.
   */
  readonly signature: string
} | {
  readonly type: "redacted_thinking"
  /**
   * Thinking content which was flagged by Amazon Bedrock's safety systems, and
   * was therefore encrypted.
   */
  readonly redactedData: string
}

declare module "@effect/ai/Prompt" {
  export interface SystemMessageOptions extends ProviderOptions {
    readonly bedrock?: {
      /**
       * Defines a section of content to be cached for reuse in subsequent API
       * calls.
       */
      readonly cachePoint?: typeof CachePointBlock.Encoded | undefined
    } | undefined
  }

  export interface UserMessageOptions extends ProviderOptions {
    readonly bedrock?: {
      /**
       * Defines a section of content to be cached for reuse in subsequent API
       * calls.
       */
      readonly cachePoint?: typeof CachePointBlock.Encoded | undefined
    } | undefined
  }

  export interface AssistantMessageOptions extends ProviderOptions {
    readonly bedrock?: {
      /**
       * Defines a section of content to be cached for reuse in subsequent API
       * calls.
       */
      readonly cachePoint?: typeof CachePointBlock.Encoded | undefined
    } | undefined
  }

  export interface ToolMessageOptions extends ProviderOptions {
    readonly bedrock?: {
      /**
       * Defines a section of content to be cached for reuse in subsequent API
       * calls.
       */
      readonly cachePoint?: typeof CachePointBlock.Encoded | undefined
    } | undefined
  }

  export interface ReasoningPartOptions extends ProviderOptions {
    readonly bedrock?: AmazonBedrockReasoningInfo | undefined
  }
}

declare module "@effect/ai/Response" {
  export interface ReasoningPartMetadata extends ProviderMetadata {
    readonly bedrock?: AmazonBedrockReasoningInfo | undefined
  }

  export interface FinishPartMetadata extends ProviderMetadata {
    readonly bedrock?: {
      readonly trace?: ConverseTrace | undefined
      readonly usage: {
        readonly cacheWriteInputTokens?: number | undefined
      }
    } | undefined
  }
}

// =============================================================================
// Amazon Bedrock Language Model
// =============================================================================

/**
 * @since 1.0.0
 * @category AiModels
 */
export const model = (
  model: (string & {}) | Model,
  config?: Omit<Config.Service, "model">
): AiModel.Model<"amazon-bedrock", LanguageModel.LanguageModel, AmazonBedrockClient> =>
  AiModel.make("amazon-bedrock", layer({ model, config }))

/**
 * @since 1.0.0
 * @category Constructors
 */
export const make = Effect.fnUntraced(function*(options: {
  readonly model: (string & {}) | Model
  readonly config?: Omit<Config.Service, "model">
}) {
  const client = yield* AmazonBedrockClient

  const makeRequest = Effect.fnUntraced(
    function*(providerOptions: LanguageModel.ProviderOptions) {
      const context = yield* Effect.context<never>()
      const config = { modelId: options.model, ...options.config, ...context.unsafeMap.get(Config.key) }
      const { messages, system } = yield* prepareMessages(providerOptions)
      const { additionalTools, betas, toolConfig } = yield* prepareTools(providerOptions, config)
      const responseFormat = providerOptions.responseFormat
      const request: typeof ConverseRequest.Encoded = {
        ...config,
        system,
        messages,
        // Handle tool configuration
        ...(responseFormat.type === "json"
          ? {
            toolConfig: {
              tools: [{
                toolSpec: {
                  name: responseFormat.objectName,
                  description: Tool.getDescriptionFromSchemaAst(responseFormat.schema.ast) ??
                    "Respond with a JSON object",
                  inputSchema: {
                    json: Tool.getJsonSchemaFromSchemaAst(responseFormat.schema.ast) as any
                  }
                }
              }],
              toolChoice: { tool: { name: responseFormat.objectName } }
            }
          }
          : Predicate.isNotUndefined(toolConfig.tools) && toolConfig.tools.length > 0
          ? { toolConfig }
          : {}),
        // Handle additional model request fields
        ...(Predicate.isNotUndefined(additionalTools)
          ? {
            additionalModelRequestFields: {
              ...config.additionalModelRequestFields,
              ...additionalTools
            }
          }
          : {})
      }
      return { betas, request }
    }
  )

  return yield* LanguageModel.make({
    generateText: Effect.fnUntraced(
      function*(options) {
        const { betas, request } = yield* makeRequest(options)
        annotateRequest(options.span, request)
        const anthropicBeta = betas.size > 0 ? Array.from(betas).join(",") : undefined
        const rawResponse = yield* client.converse({
          params: { "anthropic-beta": anthropicBeta },
          payload: request
        })
        annotateResponse(options.span, request, rawResponse)
        return yield* makeResponse(request, rawResponse, options)
      }
    ),
    streamText: Effect.fnUntraced(
      function*(options) {
        const { betas, request } = yield* makeRequest(options)
        annotateRequest(options.span, request)
        const anthropicBeta = betas.size > 0 ? Array.from(betas).join(",") : undefined
        const stream = client.converseStream({
          params: { "anthropic-beta": anthropicBeta },
          payload: request
        })
        return { request, stream }
      },
      (effect, options) =>
        effect.pipe(
          Effect.flatMap(({ request, stream }) => makeStreamResponse(request, stream, options)),
          Stream.unwrap,
          Stream.map((response) => {
            annotateStreamResponse(options.span, response)
            return response
          })
        )
    )
  })
})

/**
 * @since 1.0.0
 * @category Layers
 */
export const layer = (options: {
  readonly model: (string & {}) | Model
  readonly config?: Omit<Config.Service, "model">
}): Layer.Layer<LanguageModel.LanguageModel, never, AmazonBedrockClient> =>
  Layer.effect(LanguageModel.LanguageModel, make({ model: options.model, config: options.config }))

/**
 * @since 1.0.0
 * @category Configuration
 */
export const withConfigOverride: {
  (config: Config.Service): <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>
  <A, E, R>(self: Effect.Effect<A, E, R>, config: Config.Service): Effect.Effect<A, E, R>
} = dual<
  (config: Config.Service) => <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>,
  <A, E, R>(self: Effect.Effect<A, E, R>, config: Config.Service) => Effect.Effect<A, E, R>
>(2, (self, overrides) =>
  Effect.flatMap(
    Config.getOrUndefined,
    (config) => Effect.provideService(self, Config, { ...config, ...overrides })
  ))

// =============================================================================
// Prompt Conversion
// =============================================================================

const prepareMessages: (options: LanguageModel.ProviderOptions) => Effect.Effect<{
  readonly system: ReadonlyArray<typeof SystemContentBlock.Encoded>
  readonly messages: ReadonlyArray<typeof Message.Encoded>
}, AiError.AiError> = Effect.fnUntraced(
  function*(options) {
    const groups = groupMessages(options.prompt)

    const system: Array<typeof SystemContentBlock.Encoded> = []
    const messages: Array<typeof Message.Encoded> = []

    let documentCounter = 0
    const nextDocumentName = () => `document-${++documentCounter}`

    for (let i = 0; i < groups.length; i++) {
      const group = groups[i]
      const isLastGroup = i === groups.length - 1

      switch (group.type) {
        case "system": {
          if (messages.length > 0) {
            return yield* new AiError.MalformedInput({
              module: "AmazonBedrockLanguageModel",
              method: "prepareMessages",
              description: "Multiple system messages separated by user / assistant messages"
            })
          }
          for (const message of group.messages) {
            system.push({ text: message.content })
            if (Predicate.isNotUndefined(getCachePoint(message))) {
              system.push(BEDROCK_CACHE_POINT)
            }
          }
          break
        }

        case "user": {
          const content: Array<typeof ContentBlock.Encoded> = []

          for (const message of group.messages) {
            switch (message.role) {
              case "user": {
                for (let j = 0; j < message.content.length; j++) {
                  const part = message.content[j]

                  switch (part.type) {
                    case "text": {
                      content.push({ text: part.text })
                      break
                    }

                    case "file": {
                      if (part.data instanceof URL) {
                        // TODO(Max): support this
                        return yield* new AiError.MalformedInput({
                          module: "AmazonBedrockLanguageModel",
                          method: "prepareMessages",
                          description: "File URL inputs are not supported at this time"
                        })
                      }
                      if (part.mediaType.startsWith("image/")) {
                        content.push({
                          image: {
                            format: yield* getImageFormat(part.mediaType),
                            source: { bytes: convertToBase64(part.data) }
                          }
                        })
                      } else {
                        content.push({
                          document: {
                            format: yield* getDocumentFormat(part.mediaType),
                            name: nextDocumentName(),
                            source: { bytes: convertToBase64(part.data) }
                          }
                        })
                      }
                      break
                    }
                  }
                }
                break
              }

              case "tool": {
                for (const part of message.content) {
                  content.push({
                    toolResult: {
                      toolUseId: part.id,
                      content: [{ text: JSON.stringify(part.result) }]
                    }
                  })
                }
                break
              }
            }
          }

          messages.push({ role: "user", content })

          break
        }

        case "assistant": {
          const content: Array<typeof ContentBlock.Encoded> = []

          for (let j = 0; j < group.messages.length; j++) {
            const message = group.messages[j]
            const isLastMessage = j === group.messages.length - 1

            for (let k = 0; k < message.content.length; k++) {
              const part = message.content[k]
              const isLastPart = k === message.content.length - 1

              switch (part.type) {
                case "text": {
                  // Skip empty text blocks
                  if (part.text.trim().length === 0) {
                    break
                  }
                  content.push({
                    // Amazon Bedrock does not allow trailing whitespace in
                    // assistant content blocks
                    text: trimIfLast(isLastGroup, isLastMessage, isLastPart, part.text)
                  })
                  break
                }

                case "reasoning": {
                  const options = part.options.bedrock
                  if (Predicate.isNotUndefined(options)) {
                    if (options.type === "thinking") {
                      content.push({
                        reasoningContent: {
                          reasoningText: {
                            // Amazon Bedrock does not allow trailing whitespace in
                            // assistant content blocks
                            text: trimIfLast(isLastGroup, isLastMessage, isLastPart, part.text),
                            signature: options.signature
                          }
                        }
                      })
                    }
                    if (options.type === "redacted_thinking") {
                      content.push({
                        reasoningContent: {
                          redactedContent: options.redactedData
                        }
                      })
                    }
                  }
                  break
                }

                case "tool-call": {
                  content.push({
                    toolUse: {
                      toolUseId: part.id,
                      name: part.name,
                      input: part.params
                    }
                  })
                  break
                }
              }
            }

            if (getCachePoint(message)) {
              content.push(BEDROCK_CACHE_POINT)
            }
          }

          messages.push({ role: "assistant", content })

          break
        }
      }
    }

    return { system, messages }
  }
)

// =============================================================================
// Response Conversion
// =============================================================================

const makeResponse: (
  request: typeof ConverseRequest.Encoded,
  response: ConverseResponse,
  options: LanguageModel.ProviderOptions
) => Effect.Effect<
  Array<Response.PartEncoded>,
  never,
  IdGenerator.IdGenerator
> = Effect.fnUntraced(function*(request, response, options) {
  const parts: Array<Response.PartEncoded> = []

  parts.push({
    type: "response-metadata",
    modelId: request.modelId,
    timestamp: DateTime.formatIso(yield* DateTime.now)
  })

  for (const part of response.output.message.content) {
    switch (part.type) {
      case "text": {
        // The text parts should only be added to the response here if the
        // response format is `"text"`. If the response format is `"json"`,
        // then the text parts must instead be added to the response when a
        // tool call is received.
        if (options.responseFormat.type === "text") {
          parts.push({
            type: "text",
            text: part.text
          })
        }
        break
      }

      case "reasoningContent": {
        if (part.reasoningContent.type === "reasoning") {
          const signature = part.reasoningContent.reasoningText.signature
          parts.push({
            type: "reasoning",
            text: part.reasoningContent.reasoningText.text,
            metadata: Predicate.isNotUndefined(signature) ?
              { bedrock: { type: "thinking", signature } }
              : undefined
          })
        }
        if (part.reasoningContent.type === "redacted-reasoning") {
          parts.push({
            type: "reasoning",
            text: "",
            metadata: {
              bedrock: {
                type: "redacted_thinking",
                redactedData: part.reasoningContent.redactedContent
              }
            }
          })
        }
        break
      }

      case "toolUse": {
        // When a `"json"` response format is requested, the JSON that we need
        // will be returned by the tool call injected into the request
        if (options.responseFormat.type === "json") {
          parts.push({
            type: "text",
            text: JSON.stringify(part.toolUse.input)
          })
        } else {
          const providerTool = AnthropicTool.getProviderDefinedToolName(part.toolUse.name)
          const name = Predicate.isNotUndefined(providerTool) ? providerTool : part.toolUse.name
          const providerName = Predicate.isNotUndefined(providerTool) ? part.toolUse.name : undefined
          parts.push({
            type: "tool-call",
            id: part.toolUse.toolUseId,
            name,
            params: part.toolUse.input,
            providerName,
            providerExecuted: false
          })
        }
        break
      }
    }
  }

  const finishReason = InternalUtilities.resolveFinishReason(response.stopReason)

  parts.push({
    type: "finish",
    reason: finishReason,
    usage: {
      inputTokens: response.usage.inputTokens,
      outputTokens: response.usage.outputTokens,
      totalTokens: response.usage.totalTokens,
      cachedInputTokens: response.usage.cacheReadInputTokens
    },
    metadata: {
      bedrock: {
        trace: response.trace,
        usage: { cacheWriteInputTokens: response.usage.cacheWriteInputTokens }
      }
    }
  })

  return parts
})

const makeStreamResponse: (
  request: typeof ConverseRequest.Encoded,
  stream: Stream.Stream<ConverseResponseStreamEvent, AiError.AiError>,
  options: LanguageModel.ProviderOptions
) => Effect.Effect<
  Stream.Stream<Response.StreamPartEncoded, AiError.AiError>,
  never,
  IdGenerator.IdGenerator
> = Effect.fnUntraced(
  function*(request, stream, options) {
    const contentBlocks: Record<
      number,
      | {
        readonly type: "text"
      }
      | {
        readonly type: "reasoning"
      }
      | {
        readonly type: "tool-call"
        readonly id: string
        readonly name: string
        params: string
        readonly providerName?: string | undefined
        readonly providerExecuted: boolean
      }
    > = {}

    let trace: ConverseTrace | undefined = undefined
    let cacheWriteInputTokens: number | undefined = undefined
    const usage: Mutable<typeof Response.Usage.Encoded> = {
      inputTokens: undefined,
      outputTokens: undefined,
      totalTokens: undefined
    }

    return stream.pipe(
      Stream.mapEffect(Effect.fnUntraced(function*(event) {
        const parts: Array<Response.StreamPartEncoded> = []

        switch (event.type) {
          case "messageStart": {
            parts.push({
              type: "response-metadata",
              modelId: request.modelId,
              timestamp: DateTime.formatIso(yield* DateTime.now)
            })
            break
          }

          case "messageStop": {
            const reason = InternalUtilities.resolveFinishReason(event.messageStop.stopReason)
            parts.push({
              type: "finish",
              reason,
              usage,
              metadata: {
                bedrock: { trace, usage: { cacheWriteInputTokens } }
              }
            })

            break
          }

          case "contentBlockStart": {
            const index = event.contentBlockStart.contentBlockIndex
            const block = event.contentBlockStart
            if (Predicate.isNotUndefined(block.start.toolUse)) {
              const toolUse = block.start.toolUse
              const toolName = toolUse.name
              const providerTool = AnthropicTool.getProviderDefinedToolName(toolName)
              const name = Predicate.isNotUndefined(providerTool) ? providerTool : toolName
              const providerName = Predicate.isNotUndefined(providerTool) ? toolName : undefined

              contentBlocks[index] = {
                type: "tool-call",
                id: toolUse.toolUseId,
                name,
                params: "",
                providerName,
                providerExecuted: false
              }
              // Only emit tool param delta events for text responses
              if (options.responseFormat.type === "text") {
                parts.push({
                  type: "tool-params-start",
                  id: toolUse.toolUseId,
                  name: toolUse.name,
                  providerExecuted: false
                })
              }
            } else {
              contentBlocks[index] = { type: "text" }
              parts.push({
                type: "text-start",
                id: index.toString()
              })
            }
            break
          }

          case "contentBlockDelta": {
            const index = event.contentBlockDelta.contentBlockIndex
            const delta = event.contentBlockDelta.delta

            switch (delta.type) {
              case "text": {
                const block = contentBlocks[index]
                if (Predicate.isUndefined(block)) {
                  contentBlocks[index] = { type: "text" }
                  // Only emit text delta events for text responses
                  if (options.responseFormat.type === "text") {
                    parts.push({
                      type: "text-start",
                      id: index.toString()
                    })
                  }
                }
                // Only emit text delta events for text responses
                if (options.responseFormat.type === "text") {
                  parts.push({
                    type: "text-delta",
                    id: index.toString(),
                    delta: delta.text
                  })
                }
                break
              }

              case "reasoningContent": {
                if ("text" in delta.reasoningContent) {
                  const block = contentBlocks[index]
                  if (Predicate.isUndefined(block)) {
                    contentBlocks[index] = { type: "reasoning" }
                    parts.push({
                      type: "reasoning-start",
                      id: index.toString()
                    })
                  }
                  parts.push({
                    type: "reasoning-delta",
                    id: index.toString(),
                    delta: delta.reasoningContent.text
                  })
                } else if ("signature" in delta.reasoningContent) {
                  parts.push({
                    type: "reasoning-delta",
                    id: index.toString(),
                    delta: "",
                    metadata: {
                      bedrock: {
                        type: "thinking",
                        signature: delta.reasoningContent.signature
                      }
                    }
                  })
                } else {
                  parts.push({
                    type: "reasoning-delta",
                    id: index.toString(),
                    delta: "",
                    metadata: {
                      bedrock: {
                        type: "redacted_thinking",
                        redactedData: delta.reasoningContent.redactedContent
                      }
                    }
                  })
                }
                break
              }

              case "toolUse": {
                const block = contentBlocks[index]
                if (Predicate.isNotUndefined(block) && block.type === "tool-call") {
                  const params = delta.toolUse.input
                  // Only emit tool params delta events for text responses
                  if (options.responseFormat.type === "text") {
                    parts.push({
                      type: "tool-params-delta",
                      id: block.id,
                      delta: params
                    })
                  }
                  block.params += params
                }
                break
              }
            }
            break
          }

          case "contentBlockStop": {
            const index = event.contentBlockStop.contentBlockIndex
            const block = contentBlocks[index]
            if (Predicate.isNotUndefined(block)) {
              switch (block.type) {
                case "text": {
                  // Only emit text end events for text responses
                  if (options.responseFormat.type === "text") {
                    parts.push({
                      type: "text-end",
                      id: index.toString()
                    })
                  }
                  break
                }

                case "reasoning": {
                  parts.push({
                    type: "reasoning-end",
                    id: index.toString()
                  })
                  break
                }

                case "tool-call": {
                  // Only emit tool param events for text responses - for JSON
                  // responses, the structured output will be emitted as text
                  if (options.responseFormat.type === "text") {
                    parts.push({
                      type: "tool-params-end",
                      id: block.id
                    })

                    const toolName = block.name
                    const toolParams = block.params

                    const params = yield* Effect.try({
                      try: () => Tool.unsafeSecureJsonParse(toolParams),
                      catch: (cause) =>
                        new AiError.MalformedOutput({
                          module: "AmazonBedrockLanguageModel",
                          method: "makeStreamResponse",
                          description: "Failed to securely parse tool call parameters " +
                            `for tool '${toolName}':\nParameters: ${toolParams}`,
                          cause
                        })
                    })

                    parts.push({
                      type: "tool-call",
                      id: block.id,
                      name: toolName,
                      params,
                      providerName: block.providerName,
                      providerExecuted: block.providerExecuted
                    })
                  } else {
                    parts.push({
                      type: "text-start",
                      id: index.toString()
                    })
                    parts.push({
                      type: "text-delta",
                      id: index.toString(),
                      delta: block.params
                    })
                    parts.push({
                      type: "text-end",
                      id: index.toString()
                    })
                  }
                  break
                }
              }
              delete contentBlocks[index]
            }
            break
          }

          case "metadata": {
            usage.inputTokens = event.metadata.usage.inputTokens
            usage.outputTokens = event.metadata.usage.outputTokens
            usage.totalTokens = event.metadata.usage.totalTokens
            usage.cachedInputTokens = event.metadata.usage.cacheReadInputTokens
            if (Predicate.isNotUndefined(event.metadata.usage.cacheWriteInputTokens)) {
              cacheWriteInputTokens = event.metadata.usage.cacheWriteInputTokens
            }
            if (Predicate.isNotUndefined(event.metadata.trace)) {
              trace = event.metadata.trace
            }
            break
          }

          case "internalServerException": {
            parts.push({ type: "error", error: event.internalServerException })
            break
          }

          case "modelStreamErrorException": {
            parts.push({ type: "error", error: event.modelStreamErrorException })
            break
          }

          case "serviceUnavailableException": {
            parts.push({ type: "error", error: event.serviceUnavailableException })
            break
          }

          case "throttlingException": {
            parts.push({ type: "error", error: event.throttlingException })
            break
          }

          case "validationException": {
            parts.push({ type: "error", error: event.validationException })
            break
          }
        }

        return parts
      })),
      Stream.flattenIterables
    )
  }
)

// =============================================================================
// Tool Calling
// =============================================================================

const prepareTools: (
  options: LanguageModel.ProviderOptions,
  config: Config.Service
) => Effect.Effect<{
  readonly betas: ReadonlySet<string>
  readonly toolConfig: Partial<typeof ToolConfiguration.Encoded>
  readonly additionalTools?: Record<string, unknown> | undefined
}, AiError.AiError> = Effect.fnUntraced(function*(options, config) {
  const betas = new Set<string>()

  if (options.tools.length === 0) {
    return { toolConfig: {}, betas }
  }

  const isAnthropicModel = config.modelId!.includes("anthropic.")
  const userDefinedTools: Array<Tool.Any> = []
  const providerDefinedTools: Array<Tool.AnyProviderDefined> = []
  for (const tool of options.tools) {
    if (Tool.isUserDefined(tool)) {
      userDefinedTools.push(tool)
    } else {
      providerDefinedTools.push(tool as Tool.AnyProviderDefined)
    }
  }

  const hasAnthropicTools = isAnthropicModel && providerDefinedTools.length > 0

  let tools: Array<typeof AmazonBedrockTool.Encoded> = []
  let additionalTools: Record<string, unknown> | undefined = undefined

  // Handle Anthropic provider-defined tools for Anthropic models on Bedrock
  if (hasAnthropicTools) {
    const prepared = yield* prepareAnthropicTools(options, {})

    prepared.betas.forEach((beta) => betas.add(beta))

    if (Predicate.isNotUndefined(prepared.toolChoice)) {
      // For Anthropic tools on Bedrock, only the 'tool_choice' goes into
      // the `additionalModelRequestFields` parameter of the request, while
      // the tool definitions themselves are sent in the usual `toolConfig`
      additionalTools = {
        tool_choice: prepared.toolChoice
      }
    }

    // Handle conversion of provider-defined tools to Amazon Bedrock tool definitions
    for (const providerDefinedTool of providerDefinedTools) {
      tools.push({
        toolSpec: {
          name: providerDefinedTool.providerName,
          description: Tool.getDescription(providerDefinedTool as any),
          inputSchema: {
            json: Tool.getJsonSchema(providerDefinedTool as any) as any
          }
        }
      })
    }
  }

  // Handle conversion of user-defined tools to Amazon Bedrock tool definitions
  for (const tool of userDefinedTools) {
    tools.push({
      toolSpec: {
        name: tool.name,
        description: Tool.getDescription(tool as any),
        inputSchema: {
          json: Tool.getJsonSchema(tool as any) as any
        }
      }
    })
  }

  // Handle resolution of tool choice for **Amazon Bedrock** user-defined tools.
  // The tool choice for Anthropic provider-defined tools is resolved above and
  // inserted into the additional tool configuration object.
  let toolChoice: typeof ToolChoice.Encoded | undefined = undefined
  if (!hasAnthropicTools && tools.length > 0 && Predicate.isNotUndefined(options.toolChoice)) {
    if (options.toolChoice === "none") {
      tools.length = 0
      toolChoice = undefined
    } else if (options.toolChoice === "auto") {
      toolChoice = { auto: {} }
    } else if (options.toolChoice === "required") {
      toolChoice = { any: {} }
    } else if ("tool" in options.toolChoice) {
      toolChoice = { tool: { name: options.toolChoice.tool } }
    } else {
      const allowedTools = new Set(options.toolChoice.oneOf)
      tools = tools.filter((tool) => allowedTools.has(tool.toolSpec?.name))
      toolChoice = options.toolChoice.mode === "auto" ? { auto: {} } : { any: {} }
    }
  }

  const toolConfig: Partial<typeof ToolConfiguration.Encoded> = tools.length > 0
    ? { tools, toolChoice }
    : {}

  return { additionalTools, betas, toolConfig }
})

// =============================================================================
// Telemetry
// =============================================================================

const annotateRequest = (
  span: Span,
  request: typeof ConverseRequest.Encoded
): void => {
  addGenAIAnnotations(span, {
    system: "anthropic",
    operation: { name: "chat" },
    request: {
      model: request.modelId,
      temperature: request.inferenceConfig?.temperature,
      topP: request.inferenceConfig?.topP,
      maxTokens: request.inferenceConfig?.maxTokens,
      stopSequences: request.inferenceConfig?.stopSequences ?? []
    }
  })
}

const annotateResponse = (
  span: Span,
  request: typeof ConverseRequest.Encoded,
  response: ConverseResponse
): void => {
  addGenAIAnnotations(span, {
    response: {
      model: request.modelId,
      finishReasons: response.stopReason ? [response.stopReason] : undefined
    },
    usage: {
      inputTokens: response.usage.inputTokens,
      outputTokens: response.usage.outputTokens
    }
  })
}

const annotateStreamResponse = (span: Span, part: Response.StreamPartEncoded) => {
  if (part.type === "response-metadata") {
    addGenAIAnnotations(span, {
      response: {
        id: part.id,
        model: part.modelId
      }
    })
  }
  if (part.type === "finish") {
    addGenAIAnnotations(span, {
      response: {
        finishReasons: [part.reason]
      },
      usage: {
        inputTokens: part.usage.inputTokens,
        outputTokens: part.usage.outputTokens
      }
    })
  }
}

// =============================================================================
// Utilities
// =============================================================================

type ContentGroup = SystemMessageGroup | AssistantMessageGroup | UserMessageGroup

interface SystemMessageGroup {
  readonly type: "system"
  readonly messages: Array<Prompt.SystemMessage>
}

interface AssistantMessageGroup {
  readonly type: "assistant"
  readonly messages: Array<Prompt.AssistantMessage>
}

interface UserMessageGroup {
  readonly type: "user"
  readonly messages: Array<Prompt.ToolMessage | Prompt.UserMessage>
}

const groupMessages = (prompt: Prompt.Prompt): Array<ContentGroup> => {
  const messages: Array<ContentGroup> = []
  let current: ContentGroup | undefined = undefined
  for (const message of prompt.content) {
    switch (message.role) {
      case "system": {
        if (current?.type !== "system") {
          current = { type: "system", messages: [] }
          messages.push(current)
        }
        current.messages.push(message)
        break
      }
      case "assistant": {
        if (current?.type !== "assistant") {
          current = { type: "assistant", messages: [] }
          messages.push(current)
        }
        current.messages.push(message)
        break
      }
      case "tool":
      case "user": {
        if (current?.type !== "user") {
          current = { type: "user", messages: [] }
          messages.push(current)
        }
        current.messages.push(message)
        break
      }
    }
  }
  return messages
}

/**
 * Amazon Bedrock does not allow trailing whitespace in pre-fillled assistant
 * responses, so we trim the final text part here if it's the last message in
 * the group.
 */
const trimIfLast = (
  isLastGroup: boolean,
  isLastMessage: boolean,
  isLastPart: boolean,
  text: string
) => isLastGroup && isLastMessage && isLastPart ? text.trim() : text

const getCachePoint = (
  part:
    | Prompt.SystemMessage
    | Prompt.UserMessage
    | Prompt.AssistantMessage
    | Prompt.ToolMessage
): typeof CachePointBlock.Encoded | undefined => part.options.bedrock?.cachePoint

const convertToBase64 = (data: string | Uint8Array): string =>
  typeof data === "string" ? data : Encoding.encodeBase64(data)

const DOCUMENT_MIME_TYPES: Record<string, typeof DocumentFormat.Encoded> = {
  "application/pdf": "pdf",
  "text/csv": "csv",
  "application/msword": "doc",
  "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
  "application/vnd.ms-excel": "xls",
  "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
  "text/html": "html",
  "text/plain": "txt",
  "text/markdown": "md"
}

const getDocumentFormat: (
  mediaType: string
) => Effect.Effect<typeof DocumentFormat.Encoded, AiError.AiError> = Effect.fnUntraced(
  function*(mediaType) {
    const format = DOCUMENT_MIME_TYPES[mediaType]

    if (Predicate.isUndefined(format)) {
      return yield* new AiError.MalformedInput({
        module: "AmazonBedrockLanguageModel",
        method: "getDocumentFormat",
        description: `Unsupported document MIME type: ${mediaType} - expected ` +
          `one of: ${Object.keys(DOCUMENT_MIME_TYPES)}`
      })
    }

    return format
  }
)

const getImageFormat: (
  mediaType: string
) => Effect.Effect<typeof ImageFormat.Encoded, AiError.AiError> = Effect.fnUntraced(
  function*(mediaType) {
    const format = ImageFormat.literals.find((format) => mediaType === `image/${format}`)

    if (Predicate.isUndefined(format)) {
      return yield* new AiError.MalformedInput({
        module: "AmazonBedrockLanguageModel",
        method: "getImageFormat",
        description: `Unsupported image MIME type: ${mediaType} - expected ` +
          `one of: ${ImageFormat.literals.map((format) => `image/${format}`).join(",")}`
      })
    }

    return format
  }
)
